{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ol-xbgTvlzz2"
   },
   "source": [
    "# Plotting of DMS and Chlorophyll-A for CMAQ\n",
    "\n",
    "---\n",
    "    author: Barron H. Henderson\n",
    "    date: 2021-03-23\n",
    "    last updated: 2022-04-29\n",
    "    contributors: Brett Gantt, Jeff Willison, and Golam Sarwar\n",
    "---\n",
    "\n",
    "This notebook creates figures from DMS and CHLO from OCEAN files and the files used to create them. The inputs are often only viewed to diagnose unexpected results, as are the intermediate files. The visualization notebook will only work if the results are already available. In cloud-based systems, this may require uploading the files.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specify User Input Options\n",
    "\n",
    "* User input options are described below.\n",
    "* Most users will update `dom`'\n",
    "* `inline` is set to false to avoid saving figures in repository, set inline to show figures in addition to saving them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dom : str\n",
    "#     Name of output domain. For example, 12US1, 36US3, 108NHEMI2. This is used\n",
    "#     to name outputs and inputs.\n",
    "dom = '12US1'\n",
    "\n",
    "# inline : bool\n",
    "#     Force display of plots interactively in the notebook. Otherwise, use Agg\n",
    "inline = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CMAQ-ready Files\n",
    "ocnouttmpl = f'output/{dom}/OCEAN_*_L3m*.nc'\n",
    "# Intermediate Files\n",
    "dmsoutpath = f'dmsclimatology/{dom}/dmsconcentration.{dom}.nc'\n",
    "chlotmpl = f'chlor_a/{dom}/A*.nc'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LtBUDryDLU5P"
   },
   "source": [
    "# Import libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "1Za8quom-Jw5"
   },
   "outputs": [],
   "source": [
    "if inline:\n",
    "    %matplotlib inline\n",
    "else:\n",
    "    from matplotlib import use\n",
    "    # If Agg is not available, comment out\n",
    "    use('Agg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "1Za8quom-Jw5"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import copy\n",
    "from glob import glob\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pyproj\n",
    "\n",
    "import pycno\n",
    "import PseudoNetCDF as pnc\n",
    "\n",
    "from IPython.display import clear_output, display\n",
    "\n",
    "os.environ['IOAPI_ISPH'] = '6370000.'\n",
    "warnings.simplefilter('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(f'figs/native/', exist_ok=True)\n",
    "os.makedirs(f'figs/{dom}/', exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m6YPF65mMrW2"
   },
   "source": [
    "# Visualize DMS and CHLO in CMAQ-ready File\n",
    "\n",
    "* If working in the cloud, make sure the files have been uploaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ocnoutpaths = sorted(glob(ocnouttmpl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "gdpath = ocnoutpaths[0]\n",
    "gdf = pnc.pncopen(gdpath, format='ioapi')\n",
    "gproj = gdf.getproj(withgrid=True)\n",
    "gcno = pycno.cno(proj=gproj)\n",
    "cno = pycno.cno()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = copy(plt.get_cmap('viridis'))\n",
    "cmap.set_under('grey')\n",
    "chlonorm = plt.Normalize(vmin=1e-6, vmax=4)\n",
    "dmsnorm = plt.Normalize(vmin=1e-6, vmax=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "X_376zcqJIHP",
    "outputId": "7270d9ce-9794-4ca3-e6ae-3a172d9a4fcb"
   },
   "outputs": [],
   "source": [
    "for ocnoutpath in ocnoutpaths:\n",
    "    figpath = os.path.join('figs', dom, os.path.basename(ocnoutpath)) + '.png'\n",
    "    ocnf = pnc.pncopen(ocnoutpath, format='ioapi')\n",
    "    fig, axx = plt.subplots(1, 2, figsize=(7, 4))\n",
    "    p = axx[0].pcolormesh(ocnf.variables['CHLO'][0, 0], norm=chlonorm, cmap=cmap)\n",
    "    fig.colorbar(p, orientation='horizontal', ax=axx[0], label='CHLO mg/m3', extend='min');\n",
    "    p = axx[1].pcolormesh(ocnf.variables['DMS'][0, 0], norm=dmsnorm, cmap=cmap)\n",
    "    fig.colorbar(p, orientation='horizontal', ax=axx[1], label='DMS nM', extend='min');\n",
    "    fig.suptitle(ocnoutpath)\n",
    "    for ax in axx.ravel():\n",
    "        gcno.draw(ax=ax)\n",
    "    fig.savefig(figpath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intermediate File Visualization for Debugging\n",
    "\n",
    "* The cells below require intermediate files.\n",
    "* If you are in the cloud, these files may not be present. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4gE-JC-YLgGO"
   },
   "source": [
    "## DMS Processing Vizualization\n",
    "\n",
    "* Vizualize lat/lon DMS\n",
    "* Vizualize the gridded DMS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 643
    },
    "id": "lK7vSE_CsIts",
    "outputId": "4073e39a-0bb4-4abc-f964-9cb002d9576e"
   },
   "outputs": [],
   "source": [
    "dmsfile = pnc.pncopen('dmsclimatology/dmsconcentration.nc', format='netcdf')\n",
    "times = dmsfile.getTimes()\n",
    "lon = dmsfile.variables['longitude'][:]\n",
    "lat = dmsfile.variables['latitude'][:]\n",
    "dmsvar = dmsfile.variables['DMS']\n",
    "fig, axx = plt.subplots(4, 3, figsize=(10, 8), sharex=True, sharey=True)\n",
    "cax = fig.add_axes([.1, .1, .8, .025])\n",
    "for ai, ax in enumerate(axx.ravel()):\n",
    "    p = ax.pcolormesh(lon, lat, dmsvar[ai], norm=dmsnorm, shading='nearest')\n",
    "    fig.colorbar(p, cax=cax, orientation='horizontal')\n",
    "    ax.set_title(times[ai].strftime('%Y-%m-%d'))\n",
    "    cno.draw(ax=ax)\n",
    "\n",
    "fig.savefig('figs/native/dmsconcentration.nc.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Interpolated DMS Concentrations\n",
    "\n",
    "* Overland portions of the domain are masked before making files CMAQ-ready."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 265
    },
    "id": "MIwraJEwwqlW",
    "outputId": "b8435f25-a020-4cba-a006-a6e0b6c1c5cd"
   },
   "outputs": [],
   "source": [
    "dmsdomfile = pnc.pncopen(dmsoutpath, format='netcdf')\n",
    "dmsfigpath = os.path.join('figs', dom, os.path.basename(dmsoutpath) + '.png')\n",
    "dmsvar = dmsdomfile.variables['DMS']\n",
    "times = dmsdomfile.getTimes()\n",
    "fig, axx = plt.subplots(4, 3, figsize=(10, 10), sharex=True, sharey=True)\n",
    "cax = fig.add_axes([.1, .075, .8, .025])\n",
    "lnorm = plt.matplotlib.colors.LogNorm(vmin=1, vmax=50)\n",
    "\n",
    "for ai, ax in enumerate(axx.ravel()):\n",
    "    p = ax.pcolormesh(dmsvar[ai], norm=dmsnorm)\n",
    "    fig.colorbar(p, cax=cax, orientation='horizontal',)\n",
    "    ax.set_title(times[ai].strftime('%Y-%m-%d'))\n",
    "    gcno.draw(ax=ax)\n",
    "\n",
    "fig.savefig(dmsfigpath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4gE-JC-YLgGO"
   },
   "source": [
    "## Chlor-A Processing Vizualization\n",
    "\n",
    "* Vizualize lat/lon Chlor-A\n",
    "* Vizualize the interpolated gridded Chlor-A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 286
    },
    "id": "mMjv7s609uyG",
    "outputId": "7bd9b649-aa55-4f44-e982-b981b5ccbfb4"
   },
   "outputs": [],
   "source": [
    "chlinpaths = sorted(glob('chlor_a/A*.nc'), key=lambda x: os.path.basename(x)[5:])\n",
    "chlorfs = [pnc.pncopen(tmppath).insertDimension(time=1, before='lat', multionly=True) for tmppath in  chlinpaths]\n",
    "chlorf = chlorfs[0].stack(chlorfs[1:], 'time')\n",
    "lat = chlorf.variables['lat']\n",
    "lon = chlorf.variables['lon']\n",
    "chlorvar = chlorf.variables['chlor_a']\n",
    "\n",
    "fig, axx = plt.subplots(4, 3, figsize=(10, 8), sharex=True, sharey=True, dpi=144)\n",
    "cax = fig.add_axes([.1, .1, .8, .025])\n",
    "nthin = 1\n",
    "\n",
    "lnorm = plt.matplotlib.colors.LogNorm()\n",
    "for ai, ax in enumerate(axx.ravel()):\n",
    "    p = ax.pcolormesh(\n",
    "        lon[::nthin], lat[::nthin], chlorvar[ai, ::nthin, ::nthin],\n",
    "        norm=chlonorm, shading='nearest'\n",
    "    )\n",
    "    fig.colorbar(p, cax=cax, orientation='horizontal')\n",
    "    ax.set_title(times[ai].strftime('%Y-%m-%d'))\n",
    "    cno.draw(ax=ax)\n",
    "\n",
    "fig.savefig('figs/native/chlor_a.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Note: that overland concentrations will be masked in the CMAQ-ready files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 298
    },
    "id": "-zrH5T6gzbJU",
    "outputId": "c5aed288-f595-4f68-e479-38289564a558"
   },
   "outputs": [],
   "source": [
    "chlornewgridpaths = sorted(glob(chlotmpl), key=lambda x: os.path.basename(x)[5:])\n",
    "chlordomfs = [pnc.pncopen(tmppath).insertDimension(TSTEP=1, before='ROW', multionly=True) for tmppath in  chlornewgridpaths]\n",
    "chlordomf = chlordomfs[0].stack(chlordomfs[1:], 'TSTEP')\n",
    "chlorvar = chlordomf.variables['chlor_a']\n",
    "\n",
    "fig, axx = plt.subplots(4, 3, figsize=(10, 10), sharex=True, sharey=True)\n",
    "cax = fig.add_axes([.1, .075, .8, .025])\n",
    "\n",
    "for ai, ax in enumerate(axx.ravel()):\n",
    "    p = ax.pcolormesh(chlorvar[ai], norm=chlonorm)\n",
    "    ax.set_title(times[ai].strftime('%Y-%m-%d'))\n",
    "    gcno.draw(ax=ax)\n",
    "    \n",
    "fig.colorbar(p, cax=cax, orientation='horizontal')\n",
    "fig.savefig(f'figs/{dom}/chlor_a.{dom}.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "CMAQ_DMS_ChlorA.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "anaconda",
   "language": "python",
   "name": "anaconda"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
